﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.ServiceModel.Dispatcher;
using System.ServiceModel;
using System.ServiceModel.Description;
using System.ServiceModel.Configuration;
using System.Configuration;
using System.ServiceModel.Channels;
using WcfContrib.Extensions;
using WcfContrib.Extensions.Context;
using WcfContrib.Configuration;

namespace WcfContrib.Errors
{
    #region ErrorContextHandler
    class ErrorContextHandler : IErrorHandler
    {
        List<IErrorContextHandler> _contextHandlers;

        internal ErrorContextHandler()
        {
            _contextHandlers = new List<IErrorContextHandler>();
        }

        internal ErrorContextHandler(List<IErrorContextHandler> handlers)
        {
            _contextHandlers = handlers;
        }

        internal void AddHandlers(IEnumerable<IErrorContextHandler> handlers)
        {
            _contextHandlers.AddRange(handlers);
        }

        #region IErrorHandler Members

        bool IErrorHandler.HandleError(Exception error)
        {
            ActivationContextExtension extension = ActivationContextExtension.Current;

            if (extension != null)
            {
                if (_contextHandlers.Count > 0)
                {
                    //if (ErrorHelper.ShouldLog(error))
                    {
                        ActivationErrorContext errorContext = extension.LoadErrorContext();

                        try
                        {
                            bool isFirst = true;
                            for (int i = 0; i < _contextHandlers.Count; i++)
                            {
                                IErrorContextHandler handler = (IErrorContextHandler)_contextHandlers[i];

                                if (handler.RequiredMessageAccess == MessageAccessMode.Read)
                                {
                                    if (!isFirst)
                                    {
                                        errorContext.CreateMessageCopy();
                                    }

                                    isFirst = false;
                                }
                                else if (handler.RequiredMessageAccess == MessageAccessMode.Navigate)
                                {
                                    if (errorContext.RequestMessage.State == MessageState.Closed)
                                    {
                                        errorContext.CreateMessageCopy();
                                    }
                                }

                                handler.HandleError(error, errorContext);
                            }
                        }
                        finally
                        {
                            try
                            {
                                extension.DisposeBuffer(true);
                                errorContext.CloseMessage();
                            }
                            catch { }
                        }
                    }
                }
            }

            return false;
        }

        void IErrorHandler.ProvideFault(Exception error, System.ServiceModel.Channels.MessageVersion version, ref System.ServiceModel.Channels.Message fault)
        {
        }

        #endregion
    }
    #endregion

    #region ErrorContextHandlerBehavior
    /// <summary>
    /// A service behavior for applying error context handlers for a service
    /// </summary>
    public sealed class ErrorContextHandlerBehavior : IServiceBehavior
    {
        #region Fields
        List<string> _contextHandlerTypeNames;
        List<Type> _contextHandlerTypes;
        List<IErrorContextHandler> _contextHandlers;
        ActivationContextInitializerBehavior _contextInitBehavior;
        bool isLocked;
        ErrorContextHandler _errorHandler;
        #endregion

        #region Ctors
        private ErrorContextHandlerBehavior()
        {
            _contextHandlerTypeNames = new List<string>();
            _contextHandlerTypes = new List<Type>();
            _contextHandlers = new List<IErrorContextHandler>();
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="ErrorContextHandlerBehavior"/> class.
        /// </summary>
        /// <param name="contextHandlerTypeNames">The context handler type names.</param>
        public ErrorContextHandlerBehavior(IEnumerable<string> contextHandlerTypeNames)
            : this()
        {
            _contextHandlerTypeNames.AddRange(contextHandlerTypeNames);
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="ErrorContextHandlerBehavior"/> class.
        /// </summary>
        /// <param name="contextHandlerTypeNames">The context handler type names.</param>
        public ErrorContextHandlerBehavior(params string[] contextHandlerTypeNames)
            : this((IEnumerable<string>)contextHandlerTypeNames)
        {
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="ErrorContextHandlerBehavior"/> class.
        /// </summary>
        /// <param name="contextHandlerTypes">The context handler types.</param>
        public ErrorContextHandlerBehavior(IEnumerable<Type> contextHandlerTypes)
            : this()
        {
            _contextHandlerTypes.AddRange(contextHandlerTypes);
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="ErrorContextHandlerBehavior"/> class.
        /// </summary>
        /// <param name="contextHandlerTypes">The context handler types.</param>
        public ErrorContextHandlerBehavior(params Type[] contextHandlerTypes)
            : this((IEnumerable<Type>)contextHandlerTypes)
        {
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="ErrorContextHandlerBehavior"/> class.
        /// </summary>
        /// <param name="contextHandlers">The context handlers.</param>
        public ErrorContextHandlerBehavior(IEnumerable<IErrorContextHandler> contextHandlers)
            : this()
        {
            _contextHandlers.AddRange(contextHandlers);
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="ErrorContextHandlerBehavior"/> class.
        /// </summary>
        /// <param name="contextHandlers">The context handlers.</param>
        public ErrorContextHandlerBehavior(params IErrorContextHandler[] contextHandlers)
            : this((IEnumerable<IErrorContextHandler>)contextHandlers)
        {
        }
        #endregion

        #region Add Methods
        /// <summary>
        /// Adds the context handlers by type names.
        /// </summary>
        /// <param name="contextHandlerTypeNames">The context handler type names.</param>
        public void AddContextHandler(params string[] contextHandlerTypeNames)
        {
            AddContextHandlers((IEnumerable<string>)contextHandlerTypeNames);
        }

        /// <summary>
        /// Adds the context handlers by type names.
        /// </summary>
        /// <param name="contextHandlerTypeNames">The context handler type names.</param>
        public void AddContextHandlers(IEnumerable<string> contextHandlerTypeNames)
        {
            if (!OnHandlersAdded(contextHandlerTypeNames, null, null))
            {
                _contextHandlerTypeNames.AddRange(contextHandlerTypeNames);
            }
        }

        /// <summary>
        /// Adds the context handlers by types.
        /// </summary>
        /// <param name="contextHandlerTypes">The context handler types.</param>
        public void AddContextHandler(params Type[] contextHandlerTypes)
        {
            AddContextHandlers((IEnumerable<Type>)contextHandlerTypes);
        }

        /// <summary>
        /// Adds the context handlers by types.
        /// </summary>
        /// <param name="contextHandlerTypes">The context handler types.</param>
        public void AddContextHandlers(IEnumerable<Type> contextHandlerTypes)
        {
            if (!OnHandlersAdded(null, _contextHandlerTypes, null))
            {
                _contextHandlerTypes.AddRange(contextHandlerTypes);
            }
        }

        /// <summary>
        /// Adds the context handlers.
        /// </summary>
        /// <param name="contextHandlers">The context handlers.</param>
        public void AddContextHandler(params IErrorContextHandler[] contextHandlers)
        {
            AddContextHandlers((IEnumerable<IErrorContextHandler>)contextHandlers);
        }

        /// <summary>
        /// Adds the context handlers.
        /// </summary>
        /// <param name="contextHandlers">The context handlers.</param>
        public void AddContextHandlers(IEnumerable<IErrorContextHandler> contextHandlers)
        {
            if (!OnHandlersAdded(null, null, contextHandlers))
            {
                _contextHandlers.AddRange(contextHandlers);
            }
        }

        /// <summary>
        /// Adds the context handlers to the given service description.
        /// </summary>
        /// <param name="serviceDescription">The service description.</param>
        /// <param name="contextHandlers">The context handlers.</param>
        public static void AddContextHandler(ServiceDescription serviceDescription, params IErrorContextHandler[] contextHandlers)
        {
            AddContextHandlers(serviceDescription, (IEnumerable<IErrorContextHandler>)contextHandlers);
        }

        /// <summary>
        /// Adds the context handlers to the given service description.
        /// </summary>
        /// <param name="serviceDescription">The service description.</param>
        /// <param name="contextHandlers">The context handlers.</param>
        public static void AddContextHandlers(ServiceDescription serviceDescription, IEnumerable<IErrorContextHandler> contextHandlers)
        {
            ErrorContextHandlerBehavior b = serviceDescription.Behaviors.Find<ErrorContextHandlerBehavior>();
            if (b == null)
            {
                b = new ErrorContextHandlerBehavior(contextHandlers);

                serviceDescription.Behaviors.Add(b);
            }
            else
            {
                b.AddContextHandlers(contextHandlers);
            }
        }
        #endregion

        #region Sync
        bool OnHandlersAdded(IEnumerable<string> contextHandlerTypeNames, IEnumerable<Type> contextHandlerTypes, IEnumerable<IErrorContextHandler> contextHandlers)
        {
            if (isLocked)
            {
                IEnumerable<IErrorContextHandler> addedHandlers = contextHandlers;

                if (contextHandlerTypeNames != null)
                {
                    contextHandlerTypes = contextHandlerTypeNames.Select(h => Type.GetType(h));
                }

                if (contextHandlerTypes != null)
                {
                    contextHandlers = contextHandlerTypes.Select(h => (IErrorContextHandler)Activator.CreateInstance(h));
                }

                if (addedHandlers != null)
                {
                    _contextHandlers.AddRange(addedHandlers);
                    _errorHandler.AddHandlers(addedHandlers);

                    OnHandlersLoaded(addedHandlers);
                }

                return true;
            }

            return false;
        }

        void OnHandlersLoaded(IEnumerable<IErrorContextHandler> handlers)
        {
            bool createBuffer = handlers.Any(h => h.RequiredMessageAccess != MessageAccessMode.None);

            _contextInitBehavior.CreateBufferForErrors = createBuffer;
        }
        #endregion

        #region IServiceBehavior Members

        void IServiceBehavior.AddBindingParameters(ServiceDescription serviceDescription, ServiceHostBase serviceHostBase, System.Collections.ObjectModel.Collection<ServiceEndpoint> endpoints, System.ServiceModel.Channels.BindingParameterCollection bindingParameters)
        {
        }

        void IServiceBehavior.ApplyDispatchBehavior(ServiceDescription serviceDescription, ServiceHostBase serviceHostBase)
        {
            isLocked = true;

            _contextHandlerTypes.AddRange(_contextHandlerTypeNames.Select(name => Type.GetType(name)));
            _contextHandlers.AddRange(_contextHandlerTypes.Select(t => (IErrorContextHandler)Activator.CreateInstance(t)));

            if (_contextHandlers.Count == 0)
            {
                throw new ArgumentException("contextHandlerTypes");
            }

            _errorHandler = new ErrorContextHandler(_contextHandlers);

            foreach (ChannelDispatcher chanDisp in serviceHostBase.ChannelDispatchers)
            {
                chanDisp.ErrorHandlers.Add(_errorHandler);
            }

            _contextInitBehavior = serviceDescription.Behaviors.Find<ActivationContextInitializerBehavior>();
            if (_contextInitBehavior == null)
            {
                _contextInitBehavior = new ActivationContextInitializerBehavior();

                serviceDescription.Behaviors.Add(_contextInitBehavior);
            }

            OnHandlersLoaded(_contextHandlers);
        }

        void IServiceBehavior.Validate(ServiceDescription serviceDescription, ServiceHostBase serviceHostBase)
        {
        }

        #endregion
    }
    #endregion

    #region ErrorContextHandlerBehavior
    /// <summary>
    /// A service behavior for applying the error context handler for a service
    /// </summary>
    [AttributeUsage(AttributeTargets.Class)]
    public sealed class ErrorContextHandlerAttribute : Attribute, IServiceBehavior
    {
        List<string> _contextHandlerTypeNames;
        List<Type> _contextHandlerTypes;
        bool typeNames;

        private ErrorContextHandlerAttribute()
        {
            _contextHandlerTypeNames = new List<string>();
            _contextHandlerTypes = new List<Type>();
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="ErrorContextHandlerAttribute"/> class.
        /// </summary>
        /// <param name="contextHandlerTypeNames">The context handler type names.</param>
        public ErrorContextHandlerAttribute(params string[] contextHandlerTypeNames)
            : this()
        {
            _contextHandlerTypeNames.AddRange(contextHandlerTypeNames);
            typeNames = true;
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="ErrorContextHandlerAttribute"/> class.
        /// </summary>
        /// <param name="contextHandlerTypes">The context handler types.</param>
        public ErrorContextHandlerAttribute(params Type[] contextHandlerTypes)
            : this()
        {
            _contextHandlerTypes.AddRange(contextHandlerTypes);
        }

        #region IServiceBehavior Members

        void IServiceBehavior.AddBindingParameters(ServiceDescription serviceDescription, ServiceHostBase serviceHostBase, System.Collections.ObjectModel.Collection<ServiceEndpoint> endpoints, System.ServiceModel.Channels.BindingParameterCollection bindingParameters)
        {
        }

        void IServiceBehavior.ApplyDispatchBehavior(ServiceDescription serviceDescription, ServiceHostBase serviceHostBase)
        {
            ErrorContextHandlerBehavior errorBehavior = serviceDescription.Behaviors.Find<ErrorContextHandlerBehavior>();

            if (errorBehavior == null)
            {
                if (typeNames)
                {
                    errorBehavior = new ErrorContextHandlerBehavior(_contextHandlerTypeNames);
                }
                else
                {
                    errorBehavior = new ErrorContextHandlerBehavior(_contextHandlerTypes);
                }

                serviceDescription.Behaviors.Add(errorBehavior);
            }
            else
            {
                if (typeNames)
                {
                    errorBehavior.AddContextHandlers(_contextHandlerTypeNames);
                }
                else
                {
                    errorBehavior.AddContextHandlers(_contextHandlerTypes);
                }
            }
        }

        void IServiceBehavior.Validate(ServiceDescription serviceDescription, ServiceHostBase serviceHostBase)
        {
        }

        #endregion
    }
    #endregion

    #region ErrorContextHandlerExtensionElement
    /// <summary>
    /// Represents an error context handler behavior extension element
    /// </summary>
    public class ErrorContextHandlerExtensionElement : BehaviorExtensionElement
    {
        /// <summary>
        /// Gets or sets the error context handler elements collection.
        /// </summary>
        /// <value>The error context handler elements collection.</value>
        [ConfigurationProperty("handlers", IsDefaultCollection = true)]
        [ConfigurationCollection(typeof(TypeNameConfigurationElement), AddItemName = "handler")]
        public TypeNamesConfigurationCollection Handlers
        {
            get
            {
                return (TypeNamesConfigurationCollection)base["handlers"];
            }
            set
            {
                base["handlers"] = value;
            }
        }

        /// <summary>
        /// Gets the type of behavior.
        /// </summary>
        public override Type BehaviorType
        {
            get { return typeof(ErrorContextHandlerBehavior); }
        }

        /// <summary>
        /// Creates a behavior extension based on the current configuration settings.
        /// </summary>
        /// <returns></returns>
        protected override object CreateBehavior()
        {
            return new ErrorContextHandlerBehavior(Handlers.OfType<TypeNameConfigurationElement>().Select(t => t.TypeName));
        }
    }
    #endregion
}
